#include "VulkanApplication.h"
#include "BufferManager.h"
#include <vector>
#include <iostream>
#include <fstream>
#include <cstring>
#include <cmath>

#define STB_IMAGE_IMPLEMENTATION
#include "../glm/glm.hpp"
#include "../glm/gtc/matrix_transform.hpp"
#include "../stb_image.h"

VkImage createTextureImage(VulkanApplication& app, BufferManager& manager, const std::string& file)
{
    int w = 0, h = 0, ch = 4;
    stbi_uc* pixels = stbi_load(file.c_str(), &w, &h, &ch, STBI_rgb_alpha);
    if (!pixels)
    {
        printf("Failed to load image file.\n");
        return nullptr;
    }

    // Create image buffer
    VkDeviceSize imageSize = w * h * 4;
    VkBuffer imageBuffer = nullptr;
    VkDeviceMemory imageBufferMemory = nullptr;
    if (!manager.createBuffer(imageBuffer, imageSize, VK_BUFFER_USAGE_TRANSFER_SRC_BIT))
        return nullptr;

    VkMemoryRequirements memRequirements;
    vkGetBufferMemoryRequirements(app.getDevice(), imageBuffer, &memRequirements);
    if (!manager.createBufferMemory(imageBufferMemory, memRequirements,
                                    VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;

    void* imageData = nullptr;
    vkBindBufferMemory(app.getDevice(), imageBuffer, imageBufferMemory, 0);
    vkMapMemory(app.getDevice(), imageBufferMemory, 0, imageSize, 0, &imageData);
    memcpy(imageData, pixels, static_cast<size_t>(imageSize));
    vkUnmapMemory(app.getDevice(), imageBufferMemory);
    stbi_image_free(pixels);

    // Create image object
    VkImageCreateInfo imageInfo
    {
        VK_STRUCTURE_TYPE_IMAGE_CREATE_INFO,
        nullptr, 0,
        VK_IMAGE_TYPE_2D, VK_FORMAT_R8G8B8A8_SRGB,
        { (uint32_t)w, (uint32_t)h, 1 },
        1, 1,    // mipLevels, arrayLayers
        VK_SAMPLE_COUNT_1_BIT, VK_IMAGE_TILING_OPTIMAL,
        VK_IMAGE_USAGE_TRANSFER_DST_BIT | VK_IMAGE_USAGE_SAMPLED_BIT,
        VK_SHARING_MODE_EXCLUSIVE,
        0, nullptr,    // queueFamilyIndices
        VK_IMAGE_LAYOUT_UNDEFINED
    };

    VkImage image = nullptr;
    VkDeviceMemory imageResultBufferMemory = nullptr;
    VkResult result = vkCreateImage(app.getDevice(), &imageInfo, nullptr, &image);
    if (result != VK_SUCCESS)
    {
        printf("Failed to create image object.\n");
        return nullptr;
    }

    VkMemoryRequirements memRequirements2;
    vkGetImageMemoryRequirements(app.getDevice(), image, &memRequirements2);
    if (!manager.createBufferMemory(imageResultBufferMemory, memRequirements2,
                                    VK_MEMORY_PROPERTY_HOST_VISIBLE_BIT | VK_MEMORY_PROPERTY_HOST_COHERENT_BIT))
        return nullptr;
    else
        vkBindImageMemory(app.getDevice(), image, imageResultBufferMemory, 0);
    
    // Start copying command
    VkCommandBufferAllocateInfo allocInfo
    {
        VK_STRUCTURE_TYPE_COMMAND_BUFFER_ALLOCATE_INFO,
        nullptr, app.getCommandPool(),
        VK_COMMAND_BUFFER_LEVEL_PRIMARY, 1
    };
    VkCommandBuffer tempCommand;
    vkAllocateCommandBuffers(app.getDevice(), &allocInfo, &tempCommand);

    VkCommandBufferBeginInfo beginInfo
    {
        VK_STRUCTURE_TYPE_COMMAND_BUFFER_BEGIN_INFO,
        nullptr, VK_COMMAND_BUFFER_USAGE_ONE_TIME_SUBMIT_BIT, nullptr
    };
    vkBeginCommandBuffer(tempCommand, &beginInfo);

    // Set image layout to transfer-destination
    VkImageMemoryBarrier barrierPre
    {
        VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
        nullptr, 0, VK_ACCESS_TRANSFER_WRITE_BIT,
        VK_IMAGE_LAYOUT_UNDEFINED, VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL,
        VK_QUEUE_FAMILY_IGNORED, VK_QUEUE_FAMILY_IGNORED,
        image, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };
    vkCmdPipelineBarrier(tempCommand, VK_PIPELINE_STAGE_TOP_OF_PIPE_BIT, VK_PIPELINE_STAGE_TRANSFER_BIT,
                         0, 0, nullptr, 0, nullptr, 1, &barrierPre);

    // Copy from staging buffer to image result buffer
    VkBufferImageCopy region
    {
        0, 0, 0,
        { VK_IMAGE_ASPECT_COLOR_BIT, 0, 0, 1 },
        { 0, 0, 0 },
        { (uint32_t)w, (uint32_t)h, 1 }
    };
    vkCmdCopyBufferToImage(tempCommand, imageBuffer, image,
                           VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, 1, &region);

    // Set image layout to sampling-ready
    VkImageMemoryBarrier barrierPost
    {
        VK_STRUCTURE_TYPE_IMAGE_MEMORY_BARRIER,
        nullptr, VK_ACCESS_TRANSFER_WRITE_BIT, VK_ACCESS_SHADER_READ_BIT,
        VK_IMAGE_LAYOUT_TRANSFER_DST_OPTIMAL, VK_IMAGE_LAYOUT_SHADER_READ_ONLY_OPTIMAL,
        VK_QUEUE_FAMILY_IGNORED, VK_QUEUE_FAMILY_IGNORED,
        image, { VK_IMAGE_ASPECT_COLOR_BIT, 0, 1, 0, 1 }
    };
    vkCmdPipelineBarrier(tempCommand, VK_PIPELINE_STAGE_TRANSFER_BIT, VK_PIPELINE_STAGE_FRAGMENT_SHADER_BIT,
                         0, 0, nullptr, 0, nullptr, 1, &barrierPost);

    // Finish copying command and release resources
    VkQueue queue = nullptr;
    vkGetDeviceQueue(app.getDevice(), 0, 0, &queue);
    vkEndCommandBuffer(tempCommand);

    VkSubmitInfo submitInfo
    {
        VK_STRUCTURE_TYPE_SUBMIT_INFO,
        nullptr,
        0, nullptr, nullptr,  // waitSemaphores
        1, &tempCommand,
        0, nullptr    // signalSemaphore
    };
    vkQueueSubmit(queue, 1, &submitInfo, VK_NULL_HANDLE);
    vkQueueWaitIdle(queue);

    vkFreeCommandBuffers(app.getDevice(), app.getCommandPool(), 1, &tempCommand);
    vkFreeMemory(app.getDevice(), imageBufferMemory, nullptr);
    vkDestroyBuffer(app.getDevice(), imageBuffer, nullptr);
    return image;
}

int main(int argc, char** argv)
{
    VulkanApplication app("c13_texture", 1024, 768);
    if (!app.initialized())
    {
        printf("Failed to initialize VulkanApplication.\n");
        return 1;
    }

    // Create pipeline, vertices and indices
    BufferManager manager(&app);
    if (!manager.createDescriptorSet() ||
        !manager.prepare("draw_texture.vert", "draw_texture.frag"))
    {
        printf("Failed to create pipeline.\n");
        return 1;
    }

    // Create vertex & index buffer
    std::vector<Vertex> vertices;
    vertices.push_back({ -0.5f, -0.5f, 0.0f, 0.0f, 1.0f, 0.0f });
    vertices.push_back({ 0.5f, -0.5f, 0.0f, 0.0f, 0.0f, 1.0f });
    vertices.push_back({ 0.5f, 0.5f, 0.0f, 1.0f, 0.0f, 0.0f });
    vertices.push_back({ -0.5f, 0.5f, 0.0f, 1.0f, 1.0f, 0.0f });

    std::vector<unsigned int> indices;
    indices.push_back(0); indices.push_back(1); indices.push_back(2);
    indices.push_back(0); indices.push_back(2); indices.push_back(3);

    VkBuffer vertexBuffer = manager.createVertexBuffer("vertices", vertices);
    VkBuffer indexBuffer = manager.createIndexBuffer("indices", indices);

    // Create UBO buffer
    struct UniformBufferObject
    {
        glm::mat4x4 modelview;
        glm::mat4x4 projection;
    };

    float aspectRatio = (float)1024 / (float)768;
    glm::mat4 rotation = glm::rotate(glm::mat4(1.0f), glm::radians(30.0f), glm::vec3(0.0f, 0.0f, 1.0f));
    glm::vec3 up = glm::vec3(rotation * glm::vec4(0.0f, 1.0f, 0.0f, 1.0f));

    UniformBufferObject ubo;
    ubo.modelview = glm::lookAt(glm::vec3(0.0f, 0.0f, 2.0f), glm::vec3(), up);
    ubo.projection = glm::perspective(45.0f, aspectRatio, 0.1f, 100.0f);

    VkBuffer uboBuffer = manager.createUniformBuffer("ubo", sizeof(UniformBufferObject));
    void* uboData = manager.mapBufferMemory("ubo", sizeof(UniformBufferObject));
    memcpy(uboData, &ubo, sizeof(UniformBufferObject));
    manager.unmapBufferMemory("ubo");

    // Create texture image
    VkImage textureImage = createTextureImage(app, manager, "lingjing.jpg");

    // Main loop
    int width = 1024, height = 768;
    VkFence waitFence = app.getOrCreateFence("CommandFence");
    VkSemaphore waitNextImage = app.getOrCreateSemaphore("WaitNextImage");
    VkSemaphore waitSubmission = app.getOrCreateSemaphore("WaitSubmission");
    VkCommandBuffer commandBuffer = app.getCommandBuffer();
    while (!app.shouldClose())
    {
        vkWaitForFences(app.getDevice(), 1, &waitFence, VK_TRUE, UINT64_MAX);
        vkResetFences(app.getDevice(), 1, &waitFence);

        uint32_t imageIndex = 0;
        VkResult result = vkAcquireNextImageKHR(app.getDevice(), app.getSwapchain(), UINT64_MAX,
                                                waitNextImage, VK_NULL_HANDLE, &imageIndex);
        switch (result)
        {
        case VK_SUBOPTIMAL_KHR:
        case VK_ERROR_OUT_OF_DATE_KHR:
            imageIndex = app.recreateSwapchain(width, height, waitNextImage); break;
        case VK_SUCCESS:
            break;
        default:
            printf("Next image: %d (RESULT = %d)\n", imageIndex, result);
            app.setClosed(); break;
        }

        app.beginRendering(imageIndex, width, height);
        {
            VkBuffer vertexBuffers[] = { vertexBuffer };
            VkDeviceSize offsets[] = { 0 };
            VkDescriptorSet descrptorSet = manager.getDescriptorSet();
            vkCmdBindDescriptorSets(commandBuffer, VK_PIPELINE_BIND_POINT_GRAPHICS,
                                    app.getPipelineLayout(), 0, 1, &descrptorSet, 0, nullptr);

            vkCmdBindVertexBuffers(commandBuffer, 0, 1, vertexBuffers, offsets);
            vkCmdBindIndexBuffer(commandBuffer, indexBuffer, 0, VK_INDEX_TYPE_UINT32);
            vkCmdDrawIndexed(commandBuffer, (uint32_t)indices.size(), 1, 0, 0, 0);
        }
        app.endRendering();

        app.submitAndPresent(waitNextImage, waitSubmission, waitFence, imageIndex);
        glfwPollEvents();
    }

    vkDeviceWaitIdle(app.getDevice());
    vkDestroyImage(app.getDevice(), textureImage, nullptr);
    manager.cleanup();
    return 0;
}